#include "headfiles.h"

using namespace std;

static TNODE *LeftRotation(TNODE *root) {
    TNODE *temp = root->right;
    temp->parent = root->parent;
    if (temp->parent) {
        if (temp->parent->left == root) {
            temp->parent->left = temp;
        }
        else {
            temp->parent->right = temp;
        }
    }
    
    root->right = temp->left;
    if (temp->left) {
        temp->left->parent = root;
    }
    temp->left = root;
    root->parent = temp;
    
    root = temp;
    
    return root;
}

static TNODE *RightRotation(TNODE *root) {
    TNODE *temp = root->left;
    temp->parent = root->parent;
    if (temp->parent) {
        if (temp->parent->left == root) {
            temp->parent->left = temp;
        }
        else {
            temp->parent->right = temp;
        }
    }
    
    root->left = temp->right;
    if (temp->right) {
        temp->right->parent = root;
    }
    
    temp->right = root;
    root->parent = temp;
    
    root = temp;
    
    return root;
}

static int TreeHeight(TNODE *root) {
    if (root == NULL) {
        return 0;
    }
    
    int leftHeight = TreeHeight(root->left);
    int rightHeight = TreeHeight(root->right);
    
    return 1 + ((leftHeight > rightHeight) ? leftHeight : rightHeight);
}

static int DiffFactor(TNODE *root) {
    if (root == NULL) {
        return 0;
    }
    
    int leftHeight = TreeHeight(root->left);
    int rightHeight = TreeHeight(root->right);
    
    return leftHeight - rightHeight;
}

static TNODE * NodeBalance(TNODE *n) {
    if (n == NULL) {
        return NULL;
    }
    
    TNODE *c, *root;
    root = NULL;
    int height_diff = DiffFactor(n);
    if (height_diff == 2) {
        c = n->left;
        if (DiffFactor(c) == -1) {
            n->left = LeftRotation(c);
        }
        root = RightRotation(n);
    }
    else if (height_diff == -2) {
        c = n->right;
        if (DiffFactor(c) == 1) {
            n->right = RightRotation(c);
        }
        root = LeftRotation(n);
    }
    
    // Check whether n is the tree root
    if (root) {
        return root;
    }
    else {
        return n;
    }
}

static TNODE *AddNode(TNODE *root, int val, vector<TNODE *> &node_list) {
    TNODE *p, *cur;
    cur = root;
    while (cur) {
        p = cur;
        if (cur->val < val) {
            cur = cur->right;
        }
        else if (cur->val > val) {
            cur = cur->left;
        }
        // If this value already there, return directly
        else {
            return root;
        }
    }

    TNODE *n = new TNODE;
    if (n == NULL) {
        cout << "Allocate TNODE error in function " << __FUNCTION__ << "!" << endl;
        return NULL;
    }
    memset(n, 0, sizeof(TNODE));
    n->val = val;
    node_list.push_back(n);

    if (root == NULL) {
        return n;
    }
    
    if (p->val > val) {
        p->left = n;
    }
    else {
        p->right = n;
    }
    n->parent = p;
    
    // Rebalance the tree from p
    while (p) {
        p = NodeBalance(p);
        cur = p;
        p = p->parent;
    }
    
    return cur;
}

TNODE *createBinarySearchTree(const char *file, 
                              std::vector<TNODE *> &node_list) {
    fstream fs;
    fs.open(file, iostream::in);
    if (!fs.is_open()) {
        return NULL;
    }

    string str;
    TNODE *root = NULL;
    while (fs >> str) {
        if (str.empty()) {
            break;
        }

        int val = atoi(str.c_str());
        root = AddNode(root, val, node_list);
    }

    return root;

}

TNODE *LCABST(TNODE *root, TNODE *p, TNODE *q) {
    if (!root || !p || !q) {
        return NULL;
    }

    if (p->val < root->val &&
        q->val < root->val) {
        return LCABST(root->left, p, q);
    }
    else if (p->val > root->val &&
        q->val > root->val) {
        return LCABST(root->right, p, q);
    }
    else {
        return root;
    }
}
